#李聪 2022/3/6
#基于AID制作数据集
#引入Dataset类
from PIL import Image
from torch.utils.data import Dataset
import os

from torchvision.transforms import transforms


class AID_Dataset(Dataset):
    def __init__(self,root_dir,label_dir_list):
        self.root_dir=root_dir
        self.label_list=label_dir_list
        self.tran_totensor=transforms.ToTensor()
        self.tran_resize=transforms.Resize((224,224))
        self.label_num = len(self.label_list)
        self.img_num=[]
        for i in range(self.label_num):
            img_class_path = os.path.join(root_dir, self.label_list[i])
            img_list = os.listdir(img_class_path)
            self.img_num.append(len(img_list))
    def __getitem__(self, item):
        self.length=0
        self.label_index=0
        for i in range(self.label_num):
            self.length=self.length+self.img_num[i]
            if item<self.length:
                self.label_index=i
                break
        img_class_path=os.path.join(self.root_dir,self.label_list[self.label_index])
        img_list=os.listdir(img_class_path)
        img_name=img_list[item-self.length+self.img_num[self.label_index]]
        img_path=os.path.join(img_class_path,img_name)
        img=Image.open(img_path)
        img=self.tran_resize(img)
        img=self.tran_totensor(img)
        label=self.label_index
        return img,label
    def __len__(self):
        self.leng=0
        for i in range(len(self.img_num)):
            self.leng=self.leng+self.img_num[i]
        return self.leng
if __name__ == '__main__':
    root_dir="../autodl-tmp/AID"
    label_dir_list=os.listdir(root_dir)
    AID=AID_Dataset(root_dir,label_dir_list)
    for i in range(2):
        img,label=AID[i]
        print("*****{}******".format(i))
        print(img.shape)
        print(label)